#include "CFilt.h"

CFilt::CFilt(int irow, int icol){
    row = irow;
    col = icol;
    trans_func = Mat::zeros(row, col, CV_32FC1);
    wiener_filter = Mat::zeros(row, col, CV_32FC1);
}


void CFilt::get_wf(float k, float r){

    int nRows = row;
    int nCols = col;
    int u = 0;
    int v = 0;

    if(trans_func.isContinuous())
    {
        nCols *= row;
        nRows = 1;
        float* data = trans_func.ptr<float>(0);

        for (int p = 0; p < nCols; ++p){
            v = p / (trans_func.step/4);
            u = p % (trans_func.step/4);
            data[p] = exp(-k * pow(10, -6) * pow(pow(u-col/2, 2) + pow(v-row/2, 2), 5.0/6));
        }
    }
    else {
        for (v = 0; v < nRows; ++v)
        {
            float* data = trans_func.ptr<float>(v);
            for (u = 0; u < nCols; ++u)
            {
                data[u] = exp(-(float)k * pow(10, -6) * pow(pow(u-col/2, 2) + pow(v-row/2, 2), 5.0/6));
            }
        }
    }
    //wiener_filter = ((trans_func.mul(trans_func))/(trans_func.mul(trans_func) + (float)r*0.1)) / trans_func;
    wiener_filter = trans_func.mul(trans_func) / (trans_func.mul(trans_func) + r*0.1) / trans_func;
    //cout<< wiener_filter <<endl;
    //imshow("trans", trans_func);
}

Mat CFilt::wiener(Mat img, float k, float r, float e){

    vector<Mat> BGR;
    Mat BC, BC_dft, BC_ds, BC_inv;
    Mat GC, GC_dft, GC_ds, GC_inv;
    Mat RC, RC_dft, RC_ds, RC_inv;
    Mat img_ds;
    float er;

    this->get_wf(k, r);
    img.convertTo(img, CV_32F);
    normalize(img, img, 0, 1, CV_MINMAX);
    //normalize(img, img, 0.0, 1.0, NORM_MINMAX);
    split(img, BGR);
    BGR.back().copyTo(RC);
    BGR.pop_back();
    BGR.back().copyTo(GC);
    BGR.pop_back();
    BGR.back().copyTo(BC);
    BGR.pop_back();

    dft(BC, BC_dft, DFT_REAL_OUTPUT);
    BC_ds = BC_dft.mul(wiener_filter);
    dft(BC_ds, BC_inv, DFT_INVERSE|DFT_REAL_OUTPUT);
    normalize(BC_inv, BC_inv, 0, 1, CV_MINMAX);

    dft(GC, GC_dft, DFT_REAL_OUTPUT);
    GC_ds = GC_dft.mul(wiener_filter);
    dft(GC_ds, GC_inv, DFT_INVERSE|DFT_REAL_OUTPUT);
    normalize(GC_inv, GC_inv, 0, 1, CV_MINMAX);

    dft(RC, RC_dft, DFT_REAL_OUTPUT);
    RC_ds = RC_dft.mul(wiener_filter);
    dft(RC_ds, RC_inv, DFT_INVERSE|DFT_REAL_OUTPUT);
    normalize(RC_inv, RC_inv, 0, 1, CV_MINMAX);

    er = (e-0.1)/(150.0-0.1)*0.6 + 0.4;
    //eg = 1.7 * er ;
    //if (eg>1) eg = 1;
    //pow(GC_inv, eg, GC_inv);
    //pow(GC_inv, 0.92, GC_inv);
    pow(RC_inv, er, RC_inv);
    BGR.push_back(BC_inv);
    BGR.push_back(GC_inv);
    BGR.push_back(RC_inv);

    merge(BGR, img_ds);
    normalize(img_ds, img_ds, 0, 255, NORM_MINMAX);
    img_ds.convertTo(img_ds, CV_8U);

    return img_ds;
}

Mat CFilt::wiener(Mat img, float e){

    vector<Mat> BGR;
    Mat BC, BC_dft, BC_ds, BC_inv;
    Mat GC, GC_dft, GC_ds, GC_inv;
    Mat RC, RC_dft, RC_ds, RC_inv;
    Mat img_ds;
    float er;
    Ptr<CLAHE> clahe = createCLAHE();
    clahe->setClipLimit(3);

    er = (e-0.1)/(150.0-0.1)*0.6 + 0.4;

    img.convertTo(img, CV_32F);
    normalize(img, img, 0, 1, CV_MINMAX);
    //normalize(img, img, 0.0, 1.0, NORM_MINMAX);
    split(img, BGR);
    BGR.back().copyTo(RC);
    BGR.pop_back();
    BGR.back().copyTo(GC);
    BGR.pop_back();
    BGR.back().copyTo(BC);
    BGR.pop_back();

    dft(BC, BC_dft, DFT_REAL_OUTPUT);
    BC_ds = BC_dft.mul(wiener_filter);
    dft(BC_ds, BC_inv, DFT_INVERSE|DFT_REAL_OUTPUT);
    normalize(BC_inv, BC_inv, 0, 1, NORM_MINMAX);
    //normalize(BC_inv, BC_inv, 0, 255, NORM_MINMAX);
    //BC_ds.convertTo(BC_inv, CV_8U);
    //clahe->apply(BC_inv,BC_inv);

    dft(GC, GC_dft, DFT_REAL_OUTPUT);
    GC_ds = GC_dft.mul(wiener_filter);
    dft(GC_ds, GC_inv, DFT_INVERSE|DFT_REAL_OUTPUT);
    normalize(GC_inv, GC_inv, 0, 1, NORM_MINMAX);
    //pow(BC_inv, 0.92, BC_inv);
    //normalize(GC_inv, GC_inv, 0, 255, CV_MINMAX);
    //GC_inv.convertTo(GC_inv, CV_8U);
    //clahe->apply(GC_inv,GC_inv);

    dft(RC, RC_dft, DFT_REAL_OUTPUT);
    RC_ds = RC_dft.mul(wiener_filter);
    dft(RC_ds, RC_inv, DFT_INVERSE|DFT_REAL_OUTPUT);
    normalize(RC_inv, RC_inv, 0, 1, CV_MINMAX);
    pow(RC_inv, er, RC_inv);
    //normalize(RC_inv, RC_inv, 0, 255, NORM_MINMAX);
    //RC_ds.convertTo(RC_inv, CV_8U);
    //clahe->apply(RC_inv,RC_inv);

    normalize(BC_inv, BC_inv, 0, 255, NORM_MINMAX);
    BC_inv.convertTo(BC_inv, CV_8U);
    clahe->apply(BC_inv,BC_inv);
    normalize(GC_inv, GC_inv, 0, 255, NORM_MINMAX);
    GC_inv.convertTo(GC_inv, CV_8U);
    clahe->apply(GC_inv,GC_inv);
    normalize(RC_inv, RC_inv, 0, 255, NORM_MINMAX);
    RC_inv.convertTo(RC_inv, CV_8U);
    clahe->apply(RC_inv,RC_inv);

    BGR.push_back(BC_inv);
    BGR.push_back(GC_inv);
    BGR.push_back(RC_inv);

    merge(BGR, img_ds);
    //normalize(img_ds, img_ds, 0, 255, NORM_MINMAX);
    //img_ds.convertTo(img_ds, CV_8U);

    return img_ds;
}

Mat CFilt::clahe(Mat img){

    vector<Mat> BGR;
    Mat BC, GC, RC, img_ds;
    Ptr<CLAHE> clahe = createCLAHE();

    split(img, BGR);
    BGR.back().copyTo(RC);
    BGR.pop_back();
    BGR.back().copyTo(GC);
    BGR.pop_back();
    BGR.back().copyTo(BC);
    BGR.pop_back();

    clahe->setClipLimit(3);
    clahe->apply(BC,BC);
    clahe->setClipLimit(3);
    clahe->apply(GC,GC);
    clahe->setClipLimit(3);
    clahe->apply(RC,RC);
    BGR.push_back(BC);
    BGR.push_back(GC);
    BGR.push_back(RC);
    merge(BGR, img_ds);

    //imshow("img_ds", img_ds);
    return img_ds;
}


